import torch
import torch.nn as nn
import torch.nn.functional as F

# 深度残差模块
class NetRes(nn.Module):
    def __init__(self,res_chans=32, use_1x1conv=True):
        super().__init__()
        self.conv1 = nn.Conv2d(res_chans, res_chans, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(res_chans, res_chans, kernel_size=3, padding=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(res_chans, res_chans, kernel_size=1)
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(res_chans)
        self.bn2 = nn.BatchNorm2d(res_chans)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        if self.conv3:
            x = self.conv3(x)
        out += x
        return F.relu(out)
